import os
import time
from flask import render_template, request, flash, session, redirect, url_for, send_from_directory, \
    make_response, current_app
from ..models import User
from . import wrong
from .forms import InputForm
from .. import db


# 登录检验（用户名、密码验证）
def valid_login(username, password):
    sql = "select * from users where username='%s' and password='%s';" % (username, password)
    res = db.session.execute(sql)
    # print(sql)
    db.session.commit()
    res = list(res)
    if res:
        return True
    else:
        return False


@wrong.route('/', methods=['GET', 'POST'])
def index():
    # cookie = request.cookies.get('username_password')
    return render_template('index.html')


@wrong.route("/delete_cookie/<cookie_name>")
def delete_cookie(cookie_name):
    """
        删除cookie，通过delete_cookie()的方式，
        里面是cookie的名字
        这里的删除只是让cookie过期，并不是直接删除cookie
    """
    resp = make_response("del success")
    resp.delete_cookie(cookie_name)
    return resp


# XSS注入
@wrong.route('/XSS', methods=['GET', 'POST'])
def XSS():
    info = InputForm()
    if request.method == 'POST':
        string = request.form['string']
        return render_template('wrong/XSS.html', string=string)
    return render_template('wrong/XSS_form.html', info=info)


# CSRF
@wrong.route('/CSRF', methods=['GET', 'POST'])
def CSRF():
    if request.method == 'POST':
        if valid_login(request.form['username'], request.form['password']):
            flash("成功登录！")

            username = request.form['username']

            response = redirect(url_for('wrong.transfer'))
            response.set_cookie('username', username)

            return response
            # return render_template('success.html', cookie=cookie)
        flash('错误的用户名或密码！')
    return render_template('wrong/CSRF.html')


@wrong.route('/transfer', methods=["POST", "GET"])
def transfer():
    # 从cookie中取到用户名
    username = request.cookies.get('username', None)
    # 如果没有取到，代表没有登录
    if not username:
        time.sleep(3)
        return redirect(url_for('wrong.CSRF'))
    if request.method == "POST":
        to_account = request.form.get("to_account")
        money = request.form.get("money")
        return render_template('wrong/transfer_success.html', to_account=to_account, money=money)
        # return '转账 %s 元到 %s 成功' % (money, to_account)

    response = make_response(render_template('wrong/transfer.html'))
    return response


@wrong.route('/CSRF_attack', methods=["POST", "GET"])
def CSRF_attack():
    return render_template('wrong/CSRF_attack.html')


# HTTP_Header注入
@wrong.route('/HTTP_Header', methods=['GET', 'POST'])
def HTTP_Header():
    query = request.args.get('query')
    search_result = User.query.filter_by(username=query)
    if search_result.all():
        return render_template('wrong/HTTP_Header.html', res=search_result)
    else:
        flash('未找到相关信息')
        return render_template('wrong/HTTP_Header.html')


# 目录遍历
@wrong.route('/display', methods=['GET', 'POST'])
def _display():
    display_file_list = os.listdir(current_app.config['DISPLAY_PATH'])
    path = current_app.config['DISPLAY_PATH']
    return render_template('wrong/display.html', file_list=display_file_list, path=path)


@wrong.route('/display/<path:path>/<filename>', methods=['GET', 'POST'])
def display(path, filename):
    return send_from_directory(path, filename, as_attachment=False)


# SQL注入
@wrong.route('/SQL_inject', methods=['GET', 'POST'])
def SQL_inject():
    if request.method == 'POST':
        if valid_login(request.form['username'], request.form['password']):
            flash("成功登录！")
            # session['username'] = request.form.get('username')
            return render_template('success.html')
        else:
            flash('错误的用户名或密码！')
    return render_template('wrong/SQL_inject.html')


# 浏览可下载文件并下载
@wrong.route('/Path_traversal', methods=['GET', 'POST'])
def Path_traversal():
    download_file_list = os.listdir(current_app.config['DOWNLOAD_PATH'])
    path = current_app.config['DOWNLOAD_PATH']
    return render_template('wrong/Path_traversal_and_download.html', file_list=download_file_list, path=path)


@wrong.route('/download/<path:path>/<filename>', methods=['GET', 'POST'])
def download(path, filename):
    # return send_from_directory(current_app.config['DOWNLOAD_PATH'], filename, as_attachment=True)
    # filename = filename.encode("utf-8").decode("latin1")
    return send_from_directory(path, filename, as_attachment=True)


@wrong.route('/upload', methods=['GET', 'POST'])
def upload():
    if request.method == 'POST':
        f = request.files['file']
        if f:
            basepath = os.path.dirname(__file__)  # 当前文件所在路径
            upload_path = os.path.join(basepath, r'..\static\upload\upload_file')
            f.save(os.path.join(upload_path, f.filename))
            flash("上传图片成功!")
            return redirect(url_for('wrong.upload'))
        else:
            flash('未选择图片!')
            return redirect(url_for('wrong.upload'))
    return render_template('wrong/upload.html')


@wrong.route('/oplog', methods=['GET', 'POST'])
def oplog():
    return render_template('index.html')


# 渗透测试安全漏洞之水平越权
@wrong.route('/test', methods=['GET', 'POST'])
def test():
    if request.method == 'POST':
        if valid_login(request.form['username'], request.form['password']):
            flash("成功登录！")

            username = request.form['username']
            response = redirect(url_for('wrong.jurisdiction'))
            response.set_cookie('username', username)

            return response
            # return render_template('success.html', cookie=cookie)
        flash('错误的用户名或密码！')
    return render_template('wrong/test.html')


@wrong.route('/jurisdiction', methods=['POST', 'GET'])
def jurisdiction():
    # 从cookie中取到用户名
    username = request.cookies.get('username', None)
    cookie = 'admin'
    # 如果没有取到，代表没有登录
    if username == cookie:
        return redirect(url_for('wrong._display'))
    return redirect(url_for('wrong.test'))
